# coding = utf-8

'''
针对nii.gz文件的通用处理
'''

from collections import OrderedDict
import numpy as np
from nnUNet.nnunet.preprocessing.cropping import crop_to_nonzero,create_nonzero_mask

def test():
    import SimpleITK as sitk
    data_files = ["/datasets/nnUNet/raw_data_base/nnUNet_raw_data/Task101_Kidney/imagesTr/case_00000_0000.nii.gz"]
    seg_file = "/datasets/nnUNet/raw_data_base/nnUNet_raw_data/Task101_Kidney/labelsTr/case_00000.nii.gz"
    properties = OrderedDict()
    data_itk = [sitk.ReadImage(f) for f in data_files]

    properties["original_size_of_raw_data"] = np.array(data_itk[0].GetSize())[[2, 1, 0]]
    properties["original_spacing"] = np.array(data_itk[0].GetSpacing())[[2, 1, 0]]
    properties["list_of_data_files"] = data_files
    properties["seg_file"] = seg_file

    properties["itk_origin"] = data_itk[0].GetOrigin()
    properties["itk_spacing"] = data_itk[0].GetSpacing()
    properties["itk_direction"] = data_itk[0].GetDirection()

    data_npy = np.vstack([sitk.GetArrayFromImage(d)[None] for d in data_itk])
    seg_itk = sitk.ReadImage(seg_file)
    seg_npy = sitk.GetArrayFromImage(seg_itk)[None].astype(np.float32)
    data_npy = data_npy.astype(np.float32)

    print(properties)

    print(seg_npy.shape)
    print(data_npy.shape)



    #crop_to_nonzero(data_npy, seg_npy, nonzero_label=-1)

def crop_data_analysis():
    npz_file = "/datasets/nnUNet/raw_data_base/nnUNet_cropped_data/Task101_Kidney/case_00000.npz"
    pkl_file ="/datasets/nnUNet/raw_data_base/nnUNet_cropped_data/Task101_Kidney/case_00000.pkl"

    data = np.load(npz_file)
    print(data["data"].shape)



if __name__ == '__main__':
    crop_data_analysis()